package middleware

import (
	"net/http"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/juju/ratelimit"
)

func RateLimitMiddleware(rate int64, burst int64) gin.HandlerFunc {
	bucket := ratelimit.NewBucketWithQuantum(time.Second, rate, burst)

	return func(c *gin.Context) {
		if bucket.TakeAvailable(1) == 0 {
			c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{"error": "Too many requests"})
			return
		}

		c.Next()
	}
}
